#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys

import numpy as np
import pyzed.sl as sl
import cv2
import os
# import rospy
import torch
from std_msgs.msg import Float32
from ultralytics import YOLO
from geometry_msgs.msg import Point
# out_cat = cv2.VideoWriter("C:\\Users\\Administrator\\Desktop\\save.mp4", fourcc, 24, (352, 288), True)  # 保存位置/格式
center = (250, 250)  # 圆心坐标
radius = 20  # 半径
color = (0, 0, 255)  # 颜色（这里是蓝色）
thickness = -1  # 设置为-1，表示要填充这个圆
# 初始化ROS节点
# rospy.init_node('depth_publisher', anonymous=True)
# 创建一个发布者对象
# depth_publisher = rospy.Publisher('/depth_topic', Float32, queue_size=10)
# point_zed = rospy.Publisher('/zed_point', Point, queue_size=10)
# point_zed_person = rospy.Publisher('/zed_point_person', Point, queue_size=10)
# 2. 捕获图像


# 颜色形状检测
def ShapeAndColorDetection(img):
    contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 寻找轮廓点
    for obj in contours:
        area = cv2.contourArea(obj)  # 计算轮廓内区域的面积
        cv2.drawContours(imgContour, obj, -1, (255, 0, 0), 4)  # 绘制轮廓线
        perimeter = cv2.arcLength(obj, True)  # 计算轮廓周长
        approx = cv2.approxPolyDP(obj, 0.02 * perimeter, True)  # 获取轮廓角点坐标
        CornerNum = len(approx)  # 轮廓角点的数量
        x, y, w, h = cv2.boundingRect(approx)  # 获取坐标值和宽度、高度

        # 轮廓对象分类
        if CornerNum == 3:
            objType = "Triangle"
        elif CornerNum == 4:
            if w == h:
                objType = "Square"
            else:
                objType = "Rectangle"
        elif CornerNum > 4:
            objType = "Circle"
        else:
            objType = "N"

        # 计算颜色均值
        mask = np.zeros_like(img)
        cv2.drawContours(mask, [obj], -1, 255, -1)
        mean_color = cv2.mean(imgContour, mask=mask)[:3]
        mean_color_int = tuple(map(int, mean_color))

        cv2.rectangle(imgContour, (x, y), (x + w, y + h), (0, 0, 255), 2)  # 绘制边界框
        cv2.putText(imgContour, f"{objType}", (x + (w // 2) - 30, y + (h // 2) + 20), cv2.FONT_HERSHEY_COMPLEX,
                    0.6, (0, 0, 0), 1)  # 绘制形状
        cv2.putText(imgContour, f"RGB: {mean_color_int}", (x + (w // 2) - 40, y + (h // 2) + 40), cv2.FONT_HERSHEY_COMPLEX,
                    0.6, (0, 0, 0), 1)  # 绘制RGB值


def image_capture():
    zed = sl.Camera()

    # Set configuration parameters
    input_type = sl.InputType()
    if len(sys.argv) >= 2:
        input_type.set_from_svo_file(sys.argv[1])
    init = sl.InitParameters(input_t=input_type)
    init.camera_resolution = sl.RESOLUTION.HD720
    init.depth_mode = sl.DEPTH_MODE.PERFORMANCE
    init.coordinate_units = sl.UNIT.MILLIMETER

    # Open the camera
    err = zed.open(init)

    if err != sl.ERROR_CODE.SUCCESS:
        print(repr(err))
        zed.close()
        exit(1)
    runtime = sl.RuntimeParameters()
    resolution = zed.get_camera_information().camera_configuration.resolution
    dep = sl.Mat(resolution.width, resolution.height, sl.MAT_TYPE.U8_C4)  # 深度图
    depth = sl.Mat(resolution.width, resolution.height, sl.MAT_TYPE.U8_C4)  # 深度值

    image_zed = sl.Mat(resolution.width, resolution.height, sl.MAT_TYPE.U8_C4)
    point_cloud = sl.Mat()  # 点云数据

    # rate = rospy.Rate(100) # 10HZ

    err = zed.grab(runtime)
    if err == sl.ERROR_CODE.SUCCESS:
        # 获取图像
        # timestamp = zed.get_timestamp(sl.TIME_REFERENCE.CURRENT)  # 获取图像被捕获时的时间点
        zed.retrieve_image(image_zed, sl.VIEW.LEFT, sl.MEM.CPU, resolution) # image：容器，sl.VIEW.LEFT：内容

        # 获取深度
        # cv2.imshow("Image", img)
        zed.retrieve_measure(depth, sl.MEASURE.DEPTH, sl.MEM.CPU)  # 深度值
        zed.retrieve_image(dep, sl.VIEW.DEPTH)  # 深度图
        zed.retrieve_measure(point_cloud, sl.MEASURE.XYZBGRA, sl.MEM.CPU)
        # point_map = point_cloud.get_data()
        # depth_map = depth.get_data()
        img = image_zed.get_data()  # 转换成图像数组，便于后续的显示或者储存
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        imgGray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 转灰度图
        imgBlur = cv2.GaussianBlur(imgGray, (5, 5), 1)  # 高斯模糊
        imgCanny = cv2.Canny(imgBlur, 60, 60)  # Canny算子边缘检测
        ShapeAndColorDetection(imgCanny)  # 形状和颜色检测

        cv2.imshow("Original img", img)
        cv2.imshow("imgGray", imgGray)
        cv2.imshow("imgBlur", imgBlur)
        cv2.imshow("imgCanny", imgCanny)
        cv2.imshow("Shape and Color Detection", imgContour)

        rospy.sleep(0.01)
        cv2.waitKey(0)

    cv2.destroyAllWindows()
    zed.close()
if __name__ == '__main__':
    image_capture()

